import torch
import torchvision
import numpy as np
import torchvision.transforms as transforms
from IPython import display
from matplotlib import pyplot as plt


#  下载数据库
def down_load_data(path_dir):
    mnist_train = torchvision.datasets.FashionMNIST(root=path_dir, train=True, download=True, transform=transforms.ToTensor())
    mnist_test = torchvision.datasets.FashionMNIST(root=path_dir, train=False, download=True, transform=transforms.ToTensor())
    return mnist_train, mnist_test


#  读取小批量图集
def load_data_iter(mnist_train, mnist_test, batch_size, num_workers):
    train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=num_workers)
    return train_iter, test_iter


def init_mod(num_inputs, num_outputs):
    w = torch.tensor(np.random.normal(0, 0.01, (num_inputs, num_outputs)), dtype=torch.float)
    b = torch.zeros(num_outputs, dtype=torch.float)
    w.requires_grad_(requires_grad=True)
    b.requires_grad_(requires_grad=True)
    return w, b


def softmax(x):
    x_exp = x.exp()
    partition = x_exp.sum(dim=1, keepdim=True)
    return x_exp / partition  # 这里应用了广播机制


def net(x, w, b):
    return softmax(torch.mm(x.view((-1, num_inputs)), w) + b)


# 定义交叉熵损失
def cross_entropy(y_hat, y):
    return - torch.log(y_hat.gather(1, y.view(-1, 1)))


# 计算单样本(一个样本)的准确率
def accuracy(y_hat, y):
    return (y_hat.argmax(dim=1) == y).float().mean().item()


# 评估模型net在数据集data_iter中的准确率
def evaluate_accuracy(data_iter, net, w, b):
    acc_sum, n = 0.0, 0
    for X, y in data_iter:
        acc_sum += (net(X,w,b).argmax(dim=1) == y).float().sum().item()
        n += y.shape[0]
    return acc_sum / n


def sgd(params, lr, batch_size):
    """
        这里自动求梯度模块计算得来的梯度是一个批量样本的梯度和。
        为了和原书保持一致，这里除以了batch_size，但是应该是不用除的，因为一般用PyTorch计算loss时就默认已经
        沿batch维求了平均了。
    """
    for param in params:
        # 注意这里更改param时用的param.data，避免被`autograd`记录从而影响到梯度反向传播
        param.data -= lr * param.grad / batch_size


def train(net, train_iter, test_iter, loss, num_epochs, batch_size,
              params=None, lr=None, optimizer=None):
    """ training network """
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for x, y in train_iter:
            y_hat = net(x,params[0],params[1])
            l = loss(y_hat, y).sum()

            # 梯度清零
            if optimizer is not None:
                optimizer.zero_grad()
            elif params is not None and params[0].grad is not None:
                for param in params:
                    param.grad.data.zero_()

            l.backward()
            if optimizer is None:
                sgd(params, lr, batch_size)
            else:
                optimizer.step()  # “softmax回归的简洁实现”一节将用到

            train_l_sum += l.item()

            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()
            n += y.shape[0]
        test_acc = evaluate_accuracy(test_iter, net, params[0], params[1])
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f'
              % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))


def get_fashion_mnist_labels(labels):
    """ 函数可以将数值标签转成相应的文本标签 """
    text_labels = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat',
                   'sandal', 'shirt', 'sneaker', 'bag', 'ankle boot']
    return [text_labels[int(i)] for i in labels]


def use_svg_display():
    """Use svg format to display plot in jupyter"""
    display.set_matplotlib_formats('svg')


def show_fashion_mnist(images, labels):
    """ 该函数在一行里画出多张图像和对应标签 """
    use_svg_display()
    # 这里的_表示我们忽略（不使用）的变量
    _, figs = plt.subplots(1, len(images), figsize=(12, 12))
    for f, img, lbl in zip(figs, images, labels):
        f.imshow(img.view((28, 28)).numpy())
        f.set_title(lbl)
        f.axes.get_xaxis().set_visible(False)
        f.axes.get_yaxis().set_visible(False)
    plt.show()


print("torch.__version__:", torch.__version__)
print("torchvision.__version__:", torchvision.__version__)

batch_size = 256
num_workers = 0  # 多线程数
path_dir = '~/Datasets/FashionMNIST'
num_inputs = 784
num_outputs = 10

mnist_train, mnist_test = down_load_data(path_dir)
train_iter, test_iter = load_data_iter(mnist_train, mnist_test, batch_size, num_workers)

w, b = init_mod(num_inputs, num_outputs)
# 开始训练
num_epochs, lr = 5, 0.1
train(net, train_iter, test_iter, cross_entropy, num_epochs, batch_size, [w, b], lr)

# 预测
x, y = iter(test_iter).next()
true_labels = get_fashion_mnist_labels(y.numpy())
pred_labels = get_fashion_mnist_labels(net(x,w,b).argmax(dim=1).numpy())
titles = [true + '\n' + pred for true, pred in zip(true_labels, pred_labels)]

# 画前10个预测结果
show_fashion_mnist(x[0:9], titles[0:9])
